import cv2

# 创建CSRT跟踪器
tracker = cv2.TrackerCSRT_create()

# 打开摄像头（0通常是默认摄像头）
cap = cv2.VideoCapture(0)

# 检查摄像头是否成功打开
if not cap.isOpened():
    print("Cannot open camera")
    exit()

# 读取第一帧
ret, frame = cap.read()
if not ret:
    print("Cannot receive frame")
    cap.release()
    exit()

# 选择初始跟踪区域 (ROI)
bbox = cv2.selectROI("Select Object", frame, showCrosshair=False, fromCenter=False)
# 初始化跟踪器
tracker.init(frame, bbox)

tracking_active = True

selected_bbox_center = [(bbox[0]*2+bbox[2])/2,(bbox[1]*2+bbox[3])/2]

while True:
    # 读取新的一帧
    ret, frame = cap.read()
    if not ret:
        print("Cannot receive frame")
        break

    # 更新跟踪器
    if tracking_active:
        success, bbox = tracker.update(frame)

        # 如果跟踪成功，绘制边界框
        if success:
            x, y, w, h = [int(i) for i in bbox]
            cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
            posX = (2 * x + w) / 2
            posY = (2 * y + h) / 2
            spanX = -(posX - selected_bbox_center[0])
            spanY = posY - selected_bbox_center[1]
            info_text = "Tracking:" + str(spanX) + "," + str(spanY)
        else:
            # 跟踪失败处理
            cv2.putText(frame, "Tracking failure", (10, 80), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0, 0, 255), 2)
            info_text = "Lost"
    else:
        info_text = "Waiting for selection"

    # 显示状态信息
    cv2.putText(frame, info_text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0, 255, 0), 2)
    cv2.putText(frame, "Press 's' to select new ROI, 'q' to quit", (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.5,
                (255, 255, 255), 1)

    # 显示结果
    cv2.imshow("CSRT Tracking", frame)

    # 键盘操作
    key = cv2.waitKey(1) & 0xFF
    if key == ord('q'):  # 按 'q' 退出
        break
    elif key == ord('s'):  # 按 's' 重新选择ROI
        # 注意：重新选择需要先释放当前跟踪器并重新创建初始化，以下为简单示例
        new_bbox = cv2.selectROI("Select Object", frame, showCrosshair=False, fromCenter=False)
        selected_bbox_center = [(new_bbox[0] * 2 + new_bbox[2]) / 2, (new_bbox[1] * 2 + new_bbox[3])/2]
        tracker = cv2.TrackerCSRT_create()  # 重新创建跟踪器
        tracker.init(frame, new_bbox)
        tracking_active = True

# 释放资源
cap.release()
cv2.destroyAllWindows()
